#include <zmq.hpp>
#include <iostream>
#include <string>
#include <thread>
#include <chrono>

void worker_process(const std::string& worker_name) {
    try {
        // 创建上下文
        zmq::context_t context(1);
        
        // 创建PULL套接字
        zmq::socket_t worker(context, ZMQ_PULL);
        
        // 连接到任务分发器
        worker.connect("tcp://localhost:5555");
        
        std::cout << worker_name << " worker started" << std::endl;
        
        while (true) {
            // 接收任务
            zmq::message_t task;
            auto result = worker.recv(task, zmq::recv_flags::none);
            if (result) {
                std::string task_str = task.to_string();
                
                // 检查是否是结束信号
                if (task_str == "END") {
                    std::cout << worker_name << " received END signal, stopping..." << std::endl;
                    break;
                }
                
                // 模拟任务处理
                std::cout << worker_name << " received: " << task_str 
                         << ", processing..." << std::endl;
                std::this_thread::sleep_for(std::chrono::seconds(1));
                
                std::cout << worker_name << " completed: " << task_str << std::endl;
            }
        }
    }
    catch (const zmq::error_t& e) {
        std::cerr << worker_name << " ZMQ error: " << e.what() << std::endl;
    }
    catch (const std::exception& e) {
        std::cerr << worker_name << " error: " << e.what() << std::endl;
    }
}

int main(int argc, char* argv[]) {
    if (argc != 2) {
        std::cout << "Usage: " << argv[0] << " <worker_name(B|C|D)>" << std::endl;
        return 1;
    }
    
    std::string worker_name = argv[1];
    if (worker_name != "B" && worker_name != "C" && worker_name != "D") {
        std::cout << "Worker name must be B, C or D" << std::endl;
        return 1;
    }
    
    worker_process(worker_name);
    return 0;
}
